#include "ps.h"
#include <malloc.h>
#include <stdio.h>
#include <string.h>

#pragma pack(1)
typedef union littel_endian_size_s
{
    unsigned short int length;
    unsigned char m_byte[2];
} littel_endian_size;

typedef struct pack_start_code_s
{
    unsigned char start_code[3];
    unsigned char stream_id[1];
} pack_start_code;

typedef struct program_stream_pack_header_s
{
    pack_start_code PackStart; // 4
    unsigned char Buf[9];
    unsigned char stuffinglen;
} program_stream_pack_header;

typedef struct program_stream_map_s
{
    pack_start_code PackStart;
    littel_endian_size PackLength; //we mast do exchange
} program_stream_map;

typedef struct program_stream_e_s
{
    pack_start_code PackStart;
    littel_endian_size PackLength; //we mast do exchange
    char PackInfo1[2];
    unsigned char stuffing_length;
} program_stream_e;

// #define H264_FRAME_SIZE_MAX (1024 * 1024 * 2)

unsigned long parse_time_stamp(const unsigned char *p)
{
    unsigned long b;
    //共33位，溢出后从0开始
    unsigned long val;

    //第1个字节的第5、6、7位
    b = *p++;
    val = (b & 0x0e) << 29;

    //第2个字节的8位和第3个字节的前7位
    b = (*(p++)) << 8;
    b += *(p++);
    val += ((b & 0xfffe) << 14);

    //第4个字节的8位和第5个字节的前7位
    b = (*(p++)) << 8;
    b += *(p++);
    val += ((b & 0xfffe) >> 1);

    return val;
}

static int ProgramStreamPackHeader(char *Pack, int length, char **NextPack, int *leftlength)
{
    //printf("[%s]%x %x %x %x\n", __FUNCTION__, Pack[0], Pack[1], Pack[2], Pack[3]);
    //通过 00 00 01 ba头的第14个字节的最后3位来确定头部填充了多少字节
    program_stream_pack_header *PsHead = (program_stream_pack_header *)Pack;
    unsigned char pack_stuffing_length = PsHead->stuffinglen & '\x07';

    *leftlength = length - sizeof(program_stream_pack_header) - pack_stuffing_length; //减去头和填充的字节
    *NextPack = Pack + sizeof(program_stream_pack_header) + pack_stuffing_length;
    if (*leftlength < 4)
        return 0;

    return *leftlength;
}

static int ProgramStreamMap(char *Pack, int length, char **NextPack, int *leftlength, char **PayloadData, int *PayloadDataLen)
{
    program_stream_map *PSMPack = (program_stream_map *)Pack;

    //no payload
    *PayloadData = 0;
    *PayloadDataLen = 0;

    if ((unsigned int)length < sizeof(program_stream_map))
        return 0;

    littel_endian_size psm_length;
    psm_length.m_byte[0] = PSMPack->PackLength.m_byte[1];
    psm_length.m_byte[1] = PSMPack->PackLength.m_byte[0];

    *leftlength = length - psm_length.length - sizeof(program_stream_map);
    if (*leftlength <= 0)
        return 0;

    *NextPack = Pack + psm_length.length + sizeof(program_stream_map);

    return *leftlength;
}

static int ProgramShHead(char *Pack, int length, char **NextPack, int *leftlength, char **PayloadData, int *PayloadDataLen)
{
    program_stream_map *PSMPack = (program_stream_map *)Pack;

    //no payload
    *PayloadData = 0;
    *PayloadDataLen = 0;

    if ((unsigned int)length < sizeof(program_stream_map))
        return 0;

    littel_endian_size psm_length;
    psm_length.m_byte[0] = PSMPack->PackLength.m_byte[1];
    psm_length.m_byte[1] = PSMPack->PackLength.m_byte[0];

    *leftlength = length - psm_length.length - sizeof(program_stream_map);
    if (*leftlength <= 0)
        return 0;

    *NextPack = Pack + psm_length.length + sizeof(program_stream_map);

    return *leftlength;
}

static int Pes(char *Pack, int length, char **NextPack, int *leftlength, char **PayloadData, int *PayloadDataLen)
{
    program_stream_e *PSEPack = (program_stream_e *)Pack;

    *PayloadData = 0;
    *PayloadDataLen = 0;

    if ((unsigned int)length < sizeof(program_stream_e))
        return 0;

    littel_endian_size pse_length;
    pse_length.m_byte[0] = PSEPack->PackLength.m_byte[1];
    pse_length.m_byte[1] = PSEPack->PackLength.m_byte[0];

    *PayloadDataLen = pse_length.length - 2 - 1 - PSEPack->stuffing_length;
    if (*PayloadDataLen > 0)
        *PayloadData = Pack + sizeof(program_stream_e) + PSEPack->stuffing_length;

    *leftlength = length - pse_length.length - sizeof(pack_start_code) - sizeof(littel_endian_size);
    if (*leftlength <= 0)
        return 0;

    *NextPack = Pack + sizeof(pack_start_code) + sizeof(littel_endian_size) + pse_length.length;

    return *leftlength;
}

int GetH246FromPs(char* buffer, int length, char *h264Buffer, size_t *h264length)
{
    int leftlength = 0;
    char* NextPack = 0;

    *h264length = 0;
    unsigned char* ptr = (unsigned char*)buffer;

    if (ptr[0] == 0x00 && ptr[1] == 0x00 && ptr[2] == 0x01 && ptr[3] == 0xBA && length > 0) {

        if (ProgramStreamPackHeader(buffer, length, &NextPack, &leftlength) == 0) {
            return -1;
        }

        char* PayloadData = NULL;
        int PayloadDataLen = 0;

        while ((unsigned int)leftlength >= sizeof(pack_start_code)) {
            PayloadData = NULL;
            PayloadDataLen = 0;

            if (NextPack && NextPack[0] == '\x00' && NextPack[1] == '\x00' && NextPack[2] == '\x01' && NextPack[3] == '\xE0') {

                // //接着就是流包，说明是非i帧
                    // 测试好h264 码流中，第2个包是00 00 01 ba重复
                    // 第3个是pes包，是i帧，取出时间戳
                    // zlog_debug(h264->zlogRtp, "pts  is alive = 0x%x", 0x80 & NextPack[7]);
                    // zlog_debug(h264->zlogRtp, "dts  is alive = 0x%x", 0x40 & NextPack[7]);
                    // zlog_debug(h264->zlogRtp, "test is 0x30 = 0x%x", 0x30 & NextPack[9]);
                    // zlog_debug(h264->zlogRtp, "test is 0xc0 = 0x%x", 0x0c & NextPack[9]);
                    if ((0x80 & NextPack[7]) != 0) {
                        unsigned long pts = parse_time_stamp((unsigned char*)NextPack + 9);
                    }
                    // zlog_debug(h264->zlogRtp, "h264->pts = %ld", h264->pts);
                if (Pes(NextPack, leftlength, &NextPack, &leftlength, &PayloadData, &PayloadDataLen)) {
                    if (PayloadDataLen) {
                        if (PayloadDataLen + *h264length < H264_FRAME_SIZE_MAX) {
                            // int len = PayloadData - NextPack;
                            int len2 = PayloadData - buffer;

                            if ((PayloadDataLen + len2) > length) {
                                PayloadDataLen = length - len2;
                            }

                            memcpy(h264Buffer, PayloadData, PayloadDataLen);
                            h264Buffer += PayloadDataLen;
                            *h264length += PayloadDataLen;
                        } else {
                            printf("h264 frame size exception!! %d:%d", PayloadDataLen, *h264length);
                        }
                    }
                } else {
                    if (PayloadDataLen) {
                        if (PayloadDataLen + *h264length < H264_FRAME_SIZE_MAX) {
                            // int len = PayloadData - NextPack;
                            int len2 = PayloadData - buffer;

                            if ((PayloadDataLen + len2) > length) {
                                PayloadDataLen = length - len2;
                            }

                            memcpy(h264Buffer, PayloadData, PayloadDataLen);
                            h264Buffer += PayloadDataLen;
                            *h264length += PayloadDataLen;
                        } else {
                            printf("h264 frame size exception!! %d:%d", PayloadDataLen, *h264length);
                        }
                    }
                    break;
                }
            } else if (NextPack && NextPack[0] == '\x00' && NextPack[1] == '\x00' && NextPack[2] == '\x01' && NextPack[3] == '\xBA') {
                //重复，跳过
                if (ProgramShHead(NextPack, leftlength, &NextPack, &leftlength, &PayloadData, &PayloadDataLen) == 0){
                    break;
                }
            } else if (NextPack && NextPack[0] == '\x00' && NextPack[1] == '\x00' && NextPack[2] == '\x01' && NextPack[3] == '\xBB') {
                //重复，跳过
                if (ProgramShHead(NextPack, leftlength, &NextPack, &leftlength, &PayloadData, &PayloadDataLen) == 0){
                    break;
                }
            } else if (NextPack && NextPack[0] == '\x00' && NextPack[1] == '\x00' && NextPack[2] == '\x01' && NextPack[3] == '\xBC') {
                //第一个包是 00 00 01 BC i帧关键帧
                if (ProgramStreamMap(NextPack, leftlength, &NextPack, &leftlength, &PayloadData, &PayloadDataLen) == 0){
                    break;
                }
            } else if (NextPack && NextPack[0] == '\x00' && NextPack[1] == '\x00' && NextPack[2] == '\x01' && (NextPack[3] == '\xC0' || NextPack[3] == '\xBD')) {
                printf("audio ps frame, skip it");
                break;
            } else {
                printf("no know %x %x %x %x", NextPack[0], NextPack[1], NextPack[2], NextPack[3]);
                break;
            }
        }
    }

    return 0;
}

int GetH246FromPs(char *buffer, int length, char *h264Buffer, size_t *h264length)
{
    int leftlength = 0;
    char *NextPack = 0;
    unsigned char *ptr = (unsigned char *)buffer;

    if (ptr[0] == 0x00 && ptr[1] == 0x00 && ptr[2] == 0x01 && ptr[3] == 0xBA && length > 0)
    {

        if (ProgramStreamPackHeader(buffer, length, &NextPack, &leftlength) == 0)
        {
            return -1;
        }

        char *PayloadData = NULL;
        int PayloadDataLen = 0;

        while ((unsigned int)leftlength >= sizeof(pack_start_code))
        {
            PayloadData = NULL;
            PayloadDataLen = 0;
            // printf("head: %x %x %x %x\n", NextPack[0], NextPack[1], NextPack[2], *(NextPack + 3));
            if (NextPack && NextPack[0] == '\x00' && NextPack[1] == '\x00' && NextPack[2] == '\x01' && NextPack[3] == '\xE0')
            {
                printf("type:e0\n");
                // //接着就是流包，说明是非i帧
                // 测试好h264 码流中，第2个包是00 00 01 ba重复
                // 第3个是pes包，是i帧，取出时间戳
                // printf( "pts  is alive = 0x%x", 0x80 & NextPack[7]);
                // printf( "dts  is alive = 0x%x", 0x40 & NextPack[7]);
                // printf( "test is 0x30 = 0x%x", 0x30 & NextPack[9]);
                // printf( "test is 0xc0 = 0x%x", 0x0c & NextPack[9]);
                if ((0x80 & NextPack[7]) != 0)
                {
                    unsigned long pts = parse_time_stamp((unsigned char *)NextPack + 9);
                }
                // printf( "h264->pts = %ld", h264->pts);
                if (Pes(NextPack, leftlength, &NextPack, &leftlength, &PayloadData, &PayloadDataLen))
                {
                    if (PayloadDataLen)
                    {
                        if (PayloadDataLen + *h264length < H264_FRAME_SIZE_MAX)
                        {
                            // int len = PayloadData - NextPack;
                            int len2 = PayloadData - buffer;

                            if ((PayloadDataLen + len2) > length)
                            {
                                PayloadDataLen = length - len2;
                            }

                            memcpy(h264Buffer, PayloadData, PayloadDataLen);
                            h264Buffer += PayloadDataLen;
                            *h264length += PayloadDataLen;
                        }
                        else
                        {
                            printf("h264 frame size exception!! %d:%d", PayloadDataLen, *h264length);
                        }
                    }
                }
                else
                {
                    if (PayloadDataLen)
                    {
                        if (PayloadDataLen + *h264length < H264_FRAME_SIZE_MAX)
                        {
                            // int len = PayloadData - NextPack;
                            int len2 = PayloadData - buffer;

                            if ((PayloadDataLen + len2) > length)
                            {
                                PayloadDataLen = length - len2;
                            }

                            memcpy(h264Buffer, PayloadData, PayloadDataLen);
                            h264Buffer += PayloadDataLen;
                            *h264length += PayloadDataLen;
                        }
                        else
                        {
                            printf("h264 frame size exception!! %d:%d", PayloadDataLen, *h264length);
                        }
                    }
                    break;
                }
            }
            else if (NextPack && NextPack[0] == '\x00' && NextPack[1] == '\x00' && NextPack[2] == '\x01' && NextPack[3] == '\xBA')
            {
                printf("type:ba\n");
                //重复，跳过
                if (ProgramShHead(NextPack, leftlength, &NextPack, &leftlength, &PayloadData, &PayloadDataLen) == 0)
                {
                    break;
                }
            }
            else if (NextPack && NextPack[0] == '\x00' && NextPack[1] == '\x00' && NextPack[2] == '\x01' && NextPack[3] == '\xBB')
            {
                printf("type:bb\n");
                //重复，跳过
                if (ProgramShHead(NextPack, leftlength, &NextPack, &leftlength, &PayloadData, &PayloadDataLen) == 0)
                {
                    break;
                }
            }
            else if (NextPack && NextPack[0] == '\x00' && NextPack[1] == '\x00' && NextPack[2] == '\x01' && NextPack[3] == '\xBC')
            {
                printf("type:bc\n");
                //第一个包是 00 00 01 BC i帧关键帧
                if (ProgramStreamMap(NextPack, leftlength, &NextPack, &leftlength, &PayloadData, &PayloadDataLen) == 0)
                {
                    break;
                }
            }
            else if (NextPack && NextPack[0] == '\x00' && NextPack[1] == '\x00' && NextPack[2] == '\x01' && (NextPack[3] == '\xC0' || NextPack[3] == '\xBD'))
            {
                printf("type:C0 BD\n");
                printf("audio ps frame, skip it");
                break;
            }
            else
            {
                printf("type:other\n");
                printf("no know %x %x %x %x\n", NextPack[0], NextPack[1], NextPack[2], NextPack[3]);
                break;
            }
        }
    }

    return 0;
}
